# ----------- 新增API：获取模型列表和切换模型 -----------
from flask import send_from_directory
import glob

"""
ResNeXt 图像分类训练监控平台 - 主应用文件
重构版本，代码结构清晰，功能模块化
"""
import os
import json
import platform
from io import BytesIO
import torch

from flask import Flask, render_template, request, jsonify, send_file
from flask_socketio import SocketIO
from PIL import Image

# 导入配置
from config import Config, config

# 导入自定义模块
from models.model_utils import get_device, get_optimal_workers, get_latest_model, load_model_and_labels, cleanup_gpu_memory
from models.dataset import get_model_transform
from models.training import run_training, run_validation
from utils.database import execute_sql_and_generate_json
from utils.file_utils import validate_json_file, validate_zip_file, create_model_zip, extract_model_from_zip
from utils.system_utils import setup_macos_multiprocessing, register_cleanup_handlers, print_startup_info, get_server_settings

# ----------------- Flask 和 SocketIO 设置 -----------------
app = Flask(__name__)

# 获取运行环境
env = os.environ.get('FLASK_ENV', 'production')  # 默认为生产环境
app_config = config.get(env, config['production'])

print(f"🌍 运行环境: {env}")
print(f"🔧 调试模式: {app_config.DEBUG}")

# 根据环境变量加载对应配置
app.config.from_object(app_config)
app.config['UPLOAD_FOLDER'] = app_config.UPLOAD_FOLDER
app.config['MODEL_FOLDER'] = app_config.MODEL_FOLDER

# 在 macOS 上使用 threading 模式，避免 eventlet 的多进程问题
if platform.system() == 'Darwin':
    socketio = SocketIO(app, async_mode='threading', cors_allowed_origins="*")
else:
    socketio = SocketIO(app, async_mode='eventlet', cors_allowed_origins="*")

os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['MODEL_FOLDER'], exist_ok=True)

# --- 全局变量来追踪当前模型 ---
current_model_name = None
generated_json_filename = None  # 用于存储生成的JSON文件名

def update_current_model_name(model_name):
    """更新当前模型名称的回调函数"""
    global current_model_name
    current_model_name = model_name

# ----------------- Flask 路由 -----------------
@app.route('/')
def index():
    return render_template('index.html')

@app.route('/test')
def test():
    return send_from_directory('.', 'test.html')

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
    
    # 验证文件类型
    if not file.filename.endswith('.json'):
        return jsonify({'error': 'Invalid file type'}), 400
    
    # 验证文件格式和大小
    is_valid, error_msg = validate_json_file(file.stream)
    if not is_valid:
        return jsonify({'error': error_msg}), 400
    
    try:
        # 根据请求来源（哪个功能页面）来决定文件名
        file_purpose = request.form.get('purpose', 'training')
        if file_purpose == 'training':
            filename = "uploaded_training_data.json"
        elif file_purpose == 'validation':
            filename = "uploaded_validation_data.json"
        else:
            filename = "uploaded_data.json"

        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        return jsonify({'success': 'File uploaded successfully', 'filename': filename})
        
    except Exception as e:
        return jsonify({'error': f'上传文件时出错: {e}'}), 500

@app.route('/upload_model', methods=['POST'])
def upload_model():
    global current_model_name
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
    if not file.filename.endswith('.zip'):
        return jsonify({'error': '无效的文件类型，请上传ZIP文件。'}), 400
    
    # 验证ZIP文件
    is_valid, error_msg = validate_zip_file(file.stream)
    if not is_valid:
        return jsonify({'error': error_msg}), 400
    
    try:
        model_filename, error = extract_model_from_zip(file, app.config['MODEL_FOLDER'])
        if error:
            return jsonify({'error': error}), 400
        
        current_model_name = model_filename
        return jsonify({'success': '模型上传并设置成功', 'model_name': current_model_name})
        
    except Exception as e:
        return jsonify({'error': f'处理ZIP文件时出错: {e}'}), 500

@app.route('/download_model')
def download_model():
    global current_model_name
    if not current_model_name:
        return "没有可供下载的模型。", 404

    model_path = os.path.join(app.config['MODEL_FOLDER'], current_model_name)
    labels_filename = "labels_" + current_model_name.split('trained_model_')[1].replace('.pth', '.json')
    labels_path = os.path.join(app.config['MODEL_FOLDER'], labels_filename)

    if not os.path.exists(model_path) or not os.path.exists(labels_path):
        return "模型或标签文件丢失。", 404

    try:
        zip_data = create_model_zip(model_path, labels_path)
        zip_filename = f"model_package_{current_model_name.split('trained_model_')[1].replace('.pth', '.zip')}"
        
        return zip_data, 200, {
            'Content-Type': 'application/zip',
            'Content-Disposition': f'attachment; filename={zip_filename}'
        }
    except Exception as e:
        return f"打包模型时出错: {e}", 500

@app.route('/predict', methods=['POST'])
def predict():
    global current_model_name
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400

    if not current_model_name:
        return jsonify({'error': '未找到已训练的模型，请先训练或上传。'}), 404

    try:
        device = get_device()
        model, labels_list = load_model_and_labels(current_model_name, app.config['MODEL_FOLDER'], device)
        model.eval()

        transform = get_model_transform(is_training=False)

        image = Image.open(file.stream).convert('RGB')
        image_tensor = transform(image).unsqueeze(0).to(device)

        with torch.no_grad():
            outputs = model(image_tensor)
            _, preds = torch.max(outputs, 1)
            prediction_idx = preds.item()
        
        # 清理GPU/MPS内存
        cleanup_gpu_memory(device)
            
        predicted_label = labels_list[prediction_idx]
        return jsonify({'prediction': predicted_label})

    except Exception as e:
        return jsonify({'error': f'推理时发生错误: {e}'}), 500

@app.route('/execute_sql', methods=['POST'])
def execute_sql():
    """执行SQL查询并生成JSON"""
    global generated_json_filename
    
    try:
        data = request.get_json()
        sql_query = data.get('sql', '').strip()
        
        if not sql_query:
            return jsonify({'error': 'SQL查询语句不能为空'}), 400
        
        result_data, error, filename = execute_sql_and_generate_json(sql_query, app.config['UPLOAD_FOLDER'])
        
        if error:
            return jsonify({'error': error}), 500
        
        generated_json_filename = filename
        
        # 返回结果预览（限制前几条记录）
        preview_data = result_data[:3] if len(result_data) > 3 else result_data
        
        return jsonify({
            'success': True,
            'filename': generated_json_filename,
            'total_records': len(result_data),
            'preview': preview_data
        })
        
    except Exception as e:
        return jsonify({'error': f'执行SQL时出错: {e}'}), 500

@app.route('/download_generated_json')
def download_generated_json():
    """下载生成的JSON文件"""
    global generated_json_filename
    
    if not generated_json_filename:
        return "没有可下载的JSON文件", 404
    
    filepath = os.path.join(app.config['UPLOAD_FOLDER'], generated_json_filename)
    
    if not os.path.exists(filepath):
        return "JSON文件不存在", 404
    
    try:
        return send_file(
            filepath,
            as_attachment=True,
            download_name=generated_json_filename,
            mimetype='application/json'
        )
    except Exception as e:
        return f"下载文件时出错: {e}", 500


# 获取所有.pth模型文件列表
@app.route('/api/list_models', methods=['GET'])
def list_models():
    try:
        model_dir = app.config['MODEL_FOLDER']
        print(f"=== API: /api/list_models 被调用 ===")
        print(f"模型目录: {model_dir}")
        print(f"目录绝对路径: {os.path.abspath(model_dir)}")
        
        # 确保目录存在
        if not os.path.exists(model_dir):
            print(f"警告: 模型目录不存在: {model_dir}")
            return jsonify({'models': []})
        
        # 列出目录中的所有文件
        all_files = os.listdir(model_dir)
        print(f"目录中的所有文件: {all_files}")
        
        pattern = os.path.join(model_dir, '*.pth')
        print(f"使用模式匹配: {pattern}")
        
        model_files = [os.path.basename(f) for f in glob.glob(pattern)]
        print(f"匹配到的.pth文件: {model_files}")
        
        result = {'models': model_files}
        print(f"返回结果: {result}")
        print("=== API调用完成 ===")
        
        return jsonify(result)
    except Exception as e:
        print(f"错误: 获取模型列表失败: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

# 切换当前模型
@app.route('/api/set_model', methods=['POST'])
def set_model():
    global current_model_name
    data = request.get_json()
    model_name = data.get('model_name')
    if not model_name:
        return jsonify({'success': False, 'error': '未指定模型名'}), 400
    model_path = os.path.join(app.config['MODEL_FOLDER'], model_name)
    if not os.path.exists(model_path):
        return jsonify({'success': False, 'error': '模型文件不存在'}), 404
    current_model_name = model_name
    # current_model_name_callback(model_filename)
    return jsonify({'success': True, 'model_name': current_model_name})


# ----------------- SocketIO 事件处理 -----------------
@socketio.on('connect')
def handle_connect():
    print('客户端已连接')

@socketio.on('start_training')
def handle_start_training(data):
    filename = data.get('filename')
    if not filename:
        socketio.emit('status_update', {'msg': '错误：未提供文件名。'})
        return
        
    json_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    if not os.path.exists(json_path):
        socketio.emit('status_update', {'msg': f'错误：文件 {filename} 未找到。'})
        return
        
    # 在后台线程中启动训练
    socketio.start_background_task(
        run_training, 
        json_path, 
        app.config['MODEL_FOLDER'], 
        socketio, 
        update_current_model_name
    )

@socketio.on('start_validation')
def handle_start_validation(data):
    filename = data.get('filename')
    model_name = data.get('model_name')
    if not filename:
        socketio.emit('status_update', {'msg': '错误：未提供用于验证的文件名。'})
        return
    json_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    if not os.path.exists(json_path):
        socketio.emit('status_update', {'msg': f'错误：文件 {filename} 未找到。'})
        return
    # 优先用前端传来的模型名，否则用当前全局模型
    use_model_name = model_name if model_name else current_model_name
    # 在后台线程中启动验证
    socketio.start_background_task(
        run_validation,
        json_path,
        app.config['MODEL_FOLDER'],
        use_model_name,
        socketio
    )

# ----------------- 主程序入口 -----------------
if __name__ == '__main__':
    # --- macOS 特定设置 ---
    setup_macos_multiprocessing()

    # --- 注册信号处理和清理函数 ---
    register_cleanup_handlers()
    
    # --- 应用启动时，设置当前模型 ---
    current_model_name = get_latest_model(app.config['MODEL_FOLDER'])
    
    # --- 显示启动信息 ---
    device = get_device()
    workers_num = get_optimal_workers()
    print_startup_info(device, workers_num)
    
    try:
        # --- 启动服务器 ---
        debug_flag, use_reloader_flag = get_server_settings(env)
        
        socketio.run(app, host='0.0.0.0', port=5001, debug=debug_flag, use_reloader=use_reloader_flag)

    except KeyboardInterrupt:
        print('\n🛑 ResNeXt 训练平台服务器已退出。')
    finally:
        from utils.system_utils import cleanup_resources
        cleanup_resources()
